import torch
import os
import yaml
from tqdm import tqdm
from torch.optim import AdamW
from accelerate import Accelerator

from model_RL import build_multimodal_model # 确保这个模块和函数返回正确的模型实例
from data.balancedDataset import get_dataloaders
from transformers import get_scheduler, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import PPOTrainer, PPOConfig # 导入 PPO 相关的类
import random
from safetensors.torch import load_file as load_safetensors
import gc
import torch.nn.functional as F
import logging
import re
import copy # 用于复制模型

# --- 辅助函数 (保持不变) ---
def filter_qwen2_cache_warnings():
    """递归为所有transformers子日志器添加Qwen2缓存警告过滤器"""
    pattern = re.compile(r"Caching is incompatible with gradient checkpointing")
    class Qwen2Filter(logging.Filter):
        def filter(self, record):
            return not pattern.search(record.getMessage())
    
    def add_filter_to_loggers(logger_name):
        logger = logging.getLogger(logger_name)
        logger.addFilter(Qwen2Filter())
        for child_name in logging.Logger.manager.loggerDict:
            if child_name.startswith(logger_name + "."):
                add_filter_to_loggers(child_name)
    
    add_filter_to_loggers("transformers")

def get_dynamic_interval(step):
    if step < 1000: return 100
    elif step < 2000: return 100
    else: return 50

def get_dynamic_temperature(base_temp, step):
    decay_factor = 0.999
    return max(base_temp * (decay_factor ** step), 0.1)

def get_dynamic_threshold(step):
    if step < 2000:
        return 0.35
    else:
        return 0.32

# --- PPO 奖励函数 (修改，确保接收正确的参数) ---
def compute_rewards(model_instance, images, generated_tag_lists, true_tags, config, device): # ⚠️ 修改: 明确传入模型实例
    rewards = []
    # 确保模型处于评估模式，以防万一
    model_instance.eval() # ⚠️ 使用传入的模型实例
    with torch.no_grad(): # 奖励计算不需要梯度
        for i, (image, generated_tag_list, true_tag) in enumerate(zip(images, generated_tag_lists, true_tags)):
            # 确保 generated_tag_list 是列表 of str，例如 ['tag1', 'tag2']
            # 如果它已经是逗号分隔的字符串，则不需要 .join
            generated_tag_str = ",".join(generated_tag_list) if isinstance(generated_tag_list, list) else generated_tag_list
            
            # 奖励 1: 图像-生成标签相似度
            # model_instance.similarity_calculator 应该是一个 ChineseCLIP 实例或类似接口
            # 确保 calculate_similarity 接收单个图像 (unsqueeze(0)) 和标签字符串列表
            # 这里假定 calculate_similarity 接收 (batch_images, batch_texts)
            img_gen_sim = model_instance.similarity_calculator.calculate_similarity(
                image.unsqueeze(0), [generated_tag_str]
            ).item()
            
            # 奖励 2: 生成标签 - 真实标签相似度 (如果真实标签可用)
            # true_tag_str = true_tag 
            # gen_true_sim = model_instance.similarity_calculator.calculate_similarity(
            #     [generated_tag_str], [true_tag_str]
            # ).item()

            # 奖励 3: 长度惩罚 / 奖励 (示例)
            # 注意：generated_tag_list 应该已经是列表，如果 generate_tags 返回的是字符串，则需要先 split
            actual_tag_list = generated_tag_list if isinstance(generated_tag_list, list) else generated_tag_str.split(',')
            length_reward = 0.0
            if len(actual_tag_list) > 1:
                length_reward = 0.1 # 鼓励多个标签
            if len(generated_tag_str) > 50: # 对整个字符串长度进行惩罚
                length_reward -= 0.05 
            
            # TODO: 可以考虑添加其他奖励，例如：
            # 4. 标签重复惩罚
            # 5. 标签多样性奖励
            # 6. 特定关键词奖励/惩罚

            # 组合奖励
            total_reward = (config.reward_weight_img_gen * img_gen_sim + 
                            # config.reward_weight_gen_true * gen_true_sim +
                            config.reward_weight_length * length_reward)
            
            rewards.append(torch.tensor(total_reward, device=device))
    
    # 重新将模型设置回训练模式
    model_instance.train() # ⚠️ 使用传入的模型实例
    return torch.stack(rewards)


# --- 评估函数 (修改，确保接收正确的模型参数和处理 generate_tags 的返回值) ---
def evaluate(model, dataloader, accelerator, config): # ⚠️ 添加 config 参数
    model.eval()
    total_loss = 0
    predictions = []           # 生成的标签 (逗号分隔字符串)
    references = []            # 真实标签 (逗号分隔字符串)
    filtered_sims_list = []    # 过滤后标签的相似度（每个样本是一个列表）
    all_sims_list = []         # 所有原始标签的相似度（每个样本是一个列表）

    num_debug_samples = 5 
    printed_samples_count = 0 

    for batch in tqdm(dataloader, desc="验证", disable=not accelerator.is_main_process):
        images = batch["images"]
        # 在评估阶段使用固定的 prompts 进行生成，更接近实际应用场景
        prompts = [
            "严格根据图像内容提取工装设计相关标签(含办公空间、商业空间、酒店空间、商业餐饮等),"
            "需满足:1.优先按照风格空间场景细化,如('中式卧室双人床','北欧风婚纱摄影店','中式特色风格餐馆'),禁止泛化标签(如'客厅','房间');"
            "2. 其次包含'设计类型和功能定位'(如'科技感展厅,互动体验','轻食餐饮,社区社交空间');"
            "3. 接下来是能直接被搜索的工装元素,如('弧形玻璃隔断','透光软膜天花')"
            "4. 最后按'材质,氛围,场景适配'类别顺序排列,同一类别最多生成一个标签(如材质类别可包含'水泥漆,木饰面'),不适用的类别直接省略;"
            "5. 若图片不是场景就仅提取相关物品标签。"
            "6. 仅用英文逗号分隔,不添加任何额外内容（如解释、符号）。仅生成图像中明确存在的标签（如图像有'圆桌包厢'才生成）,不猜测未出现的内容。"
        ] * len(images)
        tags = batch["tags"] # 真实标签，用于损失计算和指标评估

        with torch.no_grad():
            # 1. 计算验证损失 (这里仍是交叉熵和对比损失)
            # MultiModalTaggingModel 的 forward 方法应在 is_training=True 时返回损失
            # 注意：这里假设 model 的 forward 方法可以处理 `images`, `prompts`, `tags` 并返回一个字典，其中包含 'loss'
            # 如果您的 MultiModalTaggingModel.forward 在 eval 模式下行为不同，可能需要调整
            # 确保传递给 model.forward 的是正确的输入格式
            
            # ⚠️ 这里需要根据 MultiModalTaggingModel 的实际 forward 签名来调用
            # 如果您的 model.forward 是 for 监督学习的，并且在 PPO 训练中不直接使用，可以考虑移除这部分
            # 或者将其调整为 PPO 的 Value Head 预测，但这超出了当前讨论范围。
            # 为了兼容性，我们假设 model.forward 仍用于计算监督学习的损失作为参考指标。
            
            # 确保 images, prompts, tags 格式正确
            # images 应该已经是处理过的 tensor
            # prompts 是 List[str]
            # tags 是 List[str]
            
            # If your MultiModalTaggingModel's forward method supports it:
            try:
                llm_output = model(images, prompts, tags=tags, is_training=True) 
                loss = llm_output['loss']
                total_loss += loss.item()
            except Exception as e:
                accelerator.print(f"[WARNING] Skipping validation loss calculation due to error: {e}")
                total_loss = float('nan') # 无法计算损失时设置为 NaN
            
            # 2. 生成标签并获取相似度 (使用 generate_tags 方法)
            # model.generate_tags 现在返回 5 个值
            generated_tags_str_list, filtered_sims, all_sims, _, _ = model.generate_tags(
                images, 
                prompts, 
                top_k=config.eval_top_k, # 从 config 中获取 top_k
                max_new_tokens=config.ppo_max_response_length, # 评估时也限制生成长度
                do_sample=False # 评估时通常不采样，使用 beam search
            ) 

            # 3. 收集结果
            # generated_tags_str_list 已经是逗号分隔的字符串列表
            predictions.extend([t.replace('，', ',').strip() for t in generated_tags_str_list])
            references.extend([t.replace('，', ',').strip() for t in tags]) # 真实标签也进行标准化
            filtered_sims_list.extend(filtered_sims)
            all_sims_list.extend(all_sims)
            
            # 4. 调试打印 (增加相似度显示)
            if accelerator.is_main_process and printed_samples_count < num_debug_samples:
                # 确保样本索引在当前批次的范围内
                sample_in_batch_idx = random.randint(0, len(images) - 1)
                sample_pred_str = generated_tags_str_list[sample_in_batch_idx]
                sample_ref_str = tags[sample_in_batch_idx] # 原始真实标签字符串
                sample_filtered_sims = filtered_sims[sample_in_batch_idx]
                sample_all_sims = all_sims[sample_in_batch_idx]
                
                print(f"\n--- 调试样本 {printed_samples_count + 1} ---")
                print(f"模型生成标签: {sample_pred_str}")
                print(f"过滤后标签的相似度: {[round(s, 4) for s in sample_filtered_sims]}")
                print(f"所有原始标签的相似度: {[round(s, 4) for s in sample_all_sims]}")
                print(f"真实标签: {sample_ref_str}")
                printed_samples_count += 1
    
    # 计算原有指标
    avg_loss = total_loss / len(dataloader) if len(dataloader) > 0 else 0.0 # 避免除以零
    total_tp, total_fp, total_fn = 0, 0, 0
    exact_match_count = 0
    
    for pred_str, ref_str in zip(predictions, references):
        pred_set = set(t.strip() for t in pred_str.split(',') if t.strip())
        ref_set = set(t.strip() for t in ref_str.split(',') if t.strip())
        
        total_tp += len(pred_set & ref_set)
        total_fp += len(pred_set - ref_set)
        total_fn += len(ref_set - pred_set)
        
        if pred_set == ref_set:
            exact_match_count += 1

    micro_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0
    micro_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0
    micro_f1 = 2 * micro_precision * micro_recall / (micro_precision + micro_recall) if (micro_precision + micro_recall) > 0 else 0
    exact_match_accuracy = exact_match_count / len(predictions) if predictions else 0

    # ---------------- 新增：相似度指标 ----------------
    all_filtered_sims = [s for sims in filtered_sims_list for s in sims]
    avg_filtered_sim = sum(all_filtered_sims) / len(all_filtered_sims) if all_filtered_sims else 0.0

    all_original_sims = [s for sims in all_sims_list for s in sims]
    avg_original_sim = sum(all_original_sims) / len(all_original_sims) if all_original_sims else 0.0

    match_sims = []
    for pred_str, ref_str, all_sims_batch in zip(predictions, references, all_sims_list): # 迭代每个样本的all_sims
        pred_tags = [t.strip() for t in pred_str.split(',') if t.strip()]
        ref_set = set(t.strip() for t in ref_str.split(',') if t.strip())
        # all_sims_batch 对应于 pred_tags 中的每个标签的相似度 (在 generate_tags 中已排序)
        # 这里需要确保 pred_tags 和 all_sims_batch 长度匹配，并且对应关系正确
        # 如果 generate_tags 已经返回了过滤并截断后的标签和对应的sims，这里直接用 filtered_sims_list 更好
        
        # 简化逻辑，直接使用 filtered_sims_list 中的相似度作为匹配相似度
        # 因为 filtered_sims_list 已经是根据过滤和 top_k 后的结果
        # 如果需要计算 "匹配上的原始标签的相似度"，则需要更复杂的映射逻辑
        # 鉴于 filtered_sims_list 和 all_sims_list 已经包含过滤信息，我们使用它们
        match_sims.extend(all_filtered_sims) # 简单地将所有过滤后的相似度加进来作为匹配相似度
    
    avg_match_sim = sum(match_sims) / len(match_sims) if match_sims else 0.0


    return avg_loss, {
        "exact_match_accuracy": exact_match_accuracy,
        "micro_precision": micro_precision,
        "micro_recall": micro_recall,
        "micro_f1": micro_f1,
        "avg_filtered_similarity": avg_filtered_sim,
        "avg_original_similarity": avg_original_sim, # 这个是所有原始生成标签的相似度
        "avg_match_similarity": avg_match_sim # 这个是过滤后的标签的平均相似度（现在等同于 avg_filtered_similarity）
    }


# --- 主运行部分 ---
if __name__ == "__main__":
    gc.collect()
    torch.cuda.empty_cache()
    
    # 1. 加载配置
    with open('Gongzhuang/config_GZ.yaml', 'r') as f:
        config_data = yaml.safe_load(f)

    class Config:
        def __init__(self, **entries):
            self.__dict__.update(entries)
            self.num_workers = entries.get('num_workers', 4)
            self.gradient_accumulation_steps = entries.get('gradient_accumulation_steps', 4)
            self.warmup_steps = entries.get('warmup_steps', 500)
            self.early_stopping_patience = entries.get('early_stopping_patience', 3)
            self.early_stopping_metric = entries.get('early_stopping_metric', 'micro_f1')
            # PPO specific configs
            self.ppo_learning_rate = entries.get('ppo_learning_rate', 1e-5)
            self.ppo_target_kl = entries.get('ppo_target_kl', 0.01)
            self.ppo_epochs = entries.get('ppo_epochs', 4) # PPO 内部的 epoch 数
            self.ppo_mini_batch_size = entries.get('ppo_mini_batch_size', 4)
            self.ppo_batch_size = entries.get('ppo_batch_size', 16) # 一个 PPO 批次中的经验数量
            self.ppo_early_stopping = entries.get('ppo_early_stopping', False)
            self.ppo_adap_kl_ctrl = entries.get('ppo_adap_kl_ctrl', True)
            self.ppo_max_response_length = entries.get('ppo_max_response_length', 64) # 生成响应的最大长度
            self.reward_weight_img_gen = entries.get('reward_weight_img_gen', 0.5) # 奖励权重
            self.reward_weight_gen_true = entries.get('reward_weight_gen_true', 0.5)
            self.reward_weight_length = entries.get('reward_weight_length', 0.2)
            self.temperature = entries.get('temperature', 0.7) # 默认温度
            self.eval_top_k = entries.get('eval_top_k', 5) # 评估时 top_k 的值
            self.log_interval = entries.get('log_interval', 10) # 打印日志的步数间隔

    config = Config(**config_data)
    os.makedirs(config.output_dir, exist_ok=True)

    # 2. 初始化 Accelerator
    accelerator = Accelerator(
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        mixed_precision="bf16",
        log_with="tensorboard",
        project_dir=config.output_dir
    )

    # 3. 构建模型、Processor 和 Tokenizer
    filter_qwen2_cache_warnings()
    # `build_multimodal_model` 应该返回 MultiModalTaggingModel 的实例
    model = build_multimodal_model(config) 
    print("------------ Model Built ------------")
    processor = model.processor
    tokenizer = model.tokenizer

    # 4. 获取数据加载器
    print("[INFO] Loading datas...")
    train_loader, val_loader, test_loader = get_dataloaders(config, processor, tokenizer)

    # 5. 初始化 PPOTrainer
    ref_model = copy.deepcopy(model)
    for param in ref_model.parameters():
        param.requires_grad = False
    
    # PPO 配置
    ppo_config = PPOConfig(
        learning_rate=config.ppo_learning_rate,
        target_kl=config.ppo_target_kl,
        ppo_epochs=config.ppo_epochs,
        mini_batch_size=config.ppo_mini_batch_size,
        batch_size=config.ppo_batch_size, 
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        early_stopping=config.ppo_early_stopping,
        adap_kl_ctrl=config.ppo_adap_kl_ctrl,
        log_with="tensorboard",
        max_grad_norm=1.0,
    )

    ppo_trainer = PPOTrainer(
        config=ppo_config,
        model=model, # MultiModalTaggingModel 实例
        ref_model=ref_model, # MultiModalTaggingModel 的冻结副本
        tokenizer=tokenizer,
        accelerator=accelerator # 将 accelerator 传递给 PPOTrainer
    )
    
    # 6. 断点续接逻辑 (需要适配 PPO 的状态保存)
    resume_checkpoint_path = os.path.join(config.output_dir, "last_checkpoint")
    start_epoch = 0
    
    # 检查是否存在 PPOTrainer 的检查点 (PPO 训练过程中的断点)
    if os.path.exists(os.path.join(resume_checkpoint_path, "ppo_trainer_state.pt")):
        accelerator.print(f"[INFO] Resuming from PPO checkpoint: {resume_checkpoint_path}")
        try:
            # 依赖 Accelerator 加载模型、优化器等状态
            accelerator.load_state(resume_checkpoint_path)
            accelerator.print("[INFO] Accelerator 状态加载完成（模型、优化器等）")
            
            # 恢复 ref_model：从加载后的策略模型复制并冻结
            ref_model = copy.deepcopy(ppo_trainer.model)
            for param in ref_model.parameters():
                param.requires_grad = False
            ppo_trainer.ref_model = ref_model 
            
            try:
                with open(os.path.join(resume_checkpoint_path, "epoch.txt"), "r") as f:
                    start_epoch = int(f.read().strip()) + 1
                accelerator.print(f"[INFO] 恢复到 epoch {start_epoch}")
            except FileNotFoundError:
                accelerator.print("[WARNING] 未找到 epoch.txt,从 epoch 0 开始")
        except Exception as e:
            accelerator.print(f"[ERROR] 加载检查点失败：{str(e)}")
            accelerator.print("[INFO] 从初始状态开始训练")
    # 如果没有找到 PPO 检查点，则尝试加载监督学习的预训练模型 (首次 PPO 训练)
    elif os.path.exists(resume_checkpoint_path) and \
          (os.path.exists(os.path.join(resume_checkpoint_path, "model.safetensors")) or \
           os.path.exists(os.path.join(resume_checkpoint_path, "pytorch_model.bin"))):
        accelerator.print(f"[INFO] Resuming from (监督学习) checkpoint: {resume_checkpoint_path}")
        try:
            # 依赖 Accelerator 加载模型、优化器等状态
            accelerator.load_state(resume_checkpoint_path)
            accelerator.print("[INFO] Accelerator 状态加载完成（模型、优化器等）")

            # 确保 ref_model 也加载了相同的权重（从加载后的策略模型复制并冻结）
            ref_model = copy.deepcopy(ppo_trainer.model) # 从加载后的策略模型复制
            for param in ref_model.parameters():
                param.requires_grad = False
            ppo_trainer.ref_model = ref_model # 更新 ppo_trainer 内部的 ref_model 引用
            
            try:
                with open(os.path.join(resume_checkpoint_path, "epoch.txt"), "r") as f:
                    start_epoch = int(f.read().strip()) + 1
                accelerator.print(f"[INFO] 恢复到 epoch {start_epoch}")
            except FileNotFoundError:
                accelerator.print("[WARNING] 未找到 epoch.txt,从 epoch 0 开始")
        
        except Exception as e:
            accelerator.print(f"[ERROR] 加载检查点失败：{str(e)}")
            accelerator.print("[INFO] 从初始状态开始训练")
    else:
        accelerator.print("[INFO] 未找到检查点，从初始状态开始训练")


    accelerator.init_trackers("multimodal_tagging_ppo") 
    # 7. PPO 训练循环
    best_val_loss = float('inf')
    best_micro_f1 = -1.0
    early_stop_counter = 0

    print("--------------------------------- PPO Training Start ---------------------------------")
    for epoch in range(start_epoch, config.max_epoch):
        # 将模型设置为训练模式
        ppo_trainer.model.train() 
        # 将参考模型设置为评估模式 (它应该总是评估模式)
        ppo_trainer.ref_model.eval()

        total_ppo_steps = 0
        total_rewards_sum = 0
        total_gen_len = 0
        
        for step, batch in enumerate(tqdm(train_loader, desc=f"PPO Epoch {epoch+1} 训练进度", disable=not accelerator.is_main_process)):
            images = batch["images"]
            # PPO 训练通常使用固定或少量变化的 prompts 来引导生成
            prompts = [
                "严格根据图像内容提取工装设计相关标签(含办公空间、商业空间、酒店空间、商业餐饮等),"
                "需满足:1.优先按照风格空间场景细化,如('中式卧室双人床','北欧风婚纱摄影店','中式特色风格餐馆'),禁止泛化标签(如'客厅','房间');"
                "2. 其次包含'设计类型和功能定位'(如'科技感展厅,互动体验','轻食餐饮,社区社交空间');"
                "3. 接下来是能直接被搜索的工装元素,如('弧形玻璃隔断','透光软膜天花')"
                "4. 最后按'材质,氛围,场景适配'类别顺序排列,同一类别最多生成一个标签(如材质类别可包含'水泥漆,木饰面'),不适用的类别直接省略;"
                "5. 若图片不是场景就仅提取相关物品标签。"
                "6. 仅用英文逗号分隔,不添加任何额外内容（如解释、符号）。仅生成图像中明确存在的标签（如图像有'圆桌包厢'才生成）,不猜测未出现的内容。"
            ] * len(images)
            true_tags = batch["tags"] # 用于奖励计算的真实标签

            # 1. 准备 `query_tensors` 和 生成 `response_tensors`
            # ⚠️ 解包 model.generate_tags 的 5 个返回值
            decoded_tag_strings, filtered_sims, all_sims, full_generated_ids, prompt_input_ids_for_ppo = \
                ppo_trainer.model.generate_tags(
                    images=images, 
                    prompts=prompts, 
                    max_new_tokens=config.ppo_max_response_length,
                    do_sample=True, # 启用采样进行 PPO 探索
                    temperature=get_dynamic_temperature(config.temperature, total_ppo_steps), # 动态温度
                    top_k=config.eval_top_k # 可以在生成时也使用 top_k
                )

            # `query_tensors` 是 PPO 算法中的 "prompt" 部分，用于计算 KL 散度
            # 它应该对应于 `prompt_input_ids_for_ppo`，即纯文本提示的 token ID
            query_tensors = prompt_input_ids_for_ppo.to(accelerator.device)
            # `response_tensors` 是 PPO 算法中的 "response" 部分，即模型实际生成的文本序列
            # 它应该从 `full_generated_ids` 中，去除多模态前缀和纯文本提示的部分
            response_tensors = []
            
            # ⚠️ 重新计算 multi_modal_prefix_len，因为 generate_tags 内部处理
            # 最佳实践是让 generate_tags 返回这个长度，或者在 PPO 训练循环外计算一次
            # 这里为了简化，我们假设 `full_generated_ids` 的开始部分是 multimodal_prefix + query_tensors
            # 并且 `query_tensors` 已经代表了原始文本提示的长度。
            # PPO Trainer 期望 response_tensors 只包含生成的"新"部分。
            # 修正：`full_generated_ids` 已经包含了多模态前缀和 prompt_input_ids_for_ppo。
            # 我们需要提取的是 `full_generated_ids` 中 **除了** `multimodal_prefix` 和 `prompt_input_ids_for_ppo` 之外的部分。
            # generate_tags 内部已经计算了 `input_token_len = multimodal_prefix_len + text_inputs.input_ids.size(1)`
            # 这里的 `text_inputs.input_ids.size(1)` 就是 `prompt_input_ids_for_ppo[i].size(0)`
            # 所以，`input_token_len` 是 `full_generated_ids` 中要截断的起始位置。
            # 为了获取正确的 response_tensors，我们需要知道每个样本的 multimodal_prefix_len
            # 在 `MultiModalTaggingModel` 的 `generate_tags` 函数中，`multimodal_prefix_len` 是固定的，
            # 并且 `text_inputs.input_ids` 的长度在填充后也是固定的。
            # 我们可以从 `full_generated_ids` 减去 `multimodal_prefix_len` 和 `query_tensors` 的长度。
            # 重新获取多模态前缀长度 (假设是固定值，或通过模型属性获取)
            multimodal_prefix_len_fixed = ppo_trainer.model.multimodal_prefix_len # ⚠️ 确保 model 有这个属性

            for i in range(len(full_generated_ids)):
                # 计算总的输入 token 长度：多模态前缀 + 纯文本提示
                current_input_token_len = multimodal_prefix_len_fixed + query_tensors[i].size(0)
                
                if full_generated_ids[i].size(0) > current_input_token_len:
                    response_part = full_generated_ids[i][current_input_token_len:]
                    # 查找 EOS token 并截断
                    eos_idx = (response_part == tokenizer.eos_token_id).nonzero(as_tuple=True)[0]
                    if eos_idx.numel() > 0:
                        response_part = response_part[:eos_idx[0]]
                else:
                    response_part = torch.tensor([], dtype=torch.long, device=accelerator.device)
                
                response_tensors.append(response_part)
            
            # 将 response_tensors 填充到相同长度
            response_tensors = torch.nn.utils.rnn.pad_sequence(
                response_tensors, batch_first=True, padding_value=tokenizer.pad_token_id
            ).to(accelerator.device)

            # 3. 计算奖励
            # `decoded_tag_strings` 是由 `model.generate_tags` 返回的 List[str] (逗号分隔的标签字符串)
            # compute_rewards 期望的是 List[List[str]]，所以需要 split
            decoded_tag_lists_for_rewards = [tags_str.split(',') for tags_str in decoded_tag_strings] # ⚠️ 转换格式
            
            # ⚠️ 将 ppo_trainer.model 传递给奖励函数，因为它包含了 similarity_calculator
            rewards = compute_rewards(ppo_trainer.model, images, decoded_tag_lists_for_rewards, true_tags, config, accelerator.device)

            # 4. 执行 PPO 优化步骤
            stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
            total_ppo_steps += 1
            total_rewards_sum += rewards.mean().item()
            total_gen_len += response_tensors.size(1) 

            # 记录日志
            ppo_trainer.log_stats(stats, batch_idx=step) 

            # 打印 PPO 相关指标
            if (step + 1) % config.log_interval == 0 and accelerator.is_main_process: # 修正为 (step + 1)
                accelerator.print(f"[PPO DEBUG] Epoch {epoch+1}, Step {step+1}, "
                                  f"Reward: {rewards.mean().item():.4f}, "
                                  f"KL: {stats['metrics/kl']:.4f}, "
                                  f"Policy Loss: {stats['losses/policy/loss']:.4f}, "
                                  f"Value Loss: {stats['losses/value/loss']:.4f}")
        
        avg_epoch_reward = total_rewards_sum / total_ppo_steps if total_ppo_steps > 0 else 0
        avg_epoch_gen_len = total_gen_len / total_ppo_steps if total_ppo_steps > 0 else 0
        accelerator.print(f"Epoch {epoch+1}/{config.max_epoch} - Avg PPO Reward: {avg_epoch_reward:.4f} - Avg Gen Length: {avg_epoch_gen_len:.2f}")
        print("---------------------------------")
        
        # 8. 评估模型 (使用 evaluate 函数)
        # 评估阶段仍然使用原始的损失计算和指标评估模型性能
        val_loss, metrics = evaluate(ppo_trainer.model, val_loader, accelerator, config) 
        current_micro_f1 = metrics['micro_f1']
        current_avg_original_similarity = metrics['avg_original_similarity']
        current_avg_match_similarity = metrics['avg_match_similarity'] 
        accelerator.print(f"Epoch {epoch+1}/{config.max_epoch} - Val Loss: {val_loss:.4f} - Micro F1: {current_micro_f1:.4f} - avg_original_similarity: {current_avg_original_similarity:.4f} - avg_match_similarity: {current_avg_match_similarity:.4f}")

        # 9. 早停逻辑
        is_improved = False
        if config.early_stopping_metric == 'micro_f1':
            if current_micro_f1 > best_micro_f1:
                is_improved = True
        else: # 默认为 val_loss
            if val_loss < best_val_loss:
                is_improved = True
        
        if is_improved:
            early_stop_counter = 0
            # 更新最佳指标
            if config.early_stopping_metric == 'micro_f1':
                best_micro_f1 = current_micro_f1
                accelerator.print(f"Saving best model (by F1) with Micro F1: {best_micro_f1:.4f}")
            else:
                best_val_loss = val_loss
                accelerator.print(f"Saving best model (by Val Loss) with Val Loss: {best_val_loss:.4f}")
            
            # 保存最佳模型状态 (PPO Trainer 的保存方式)
            # PPOTrainer 提供了 save_pretrained 方法来保存模型和配置
            ppo_trainer.save_pretrained(os.path.join(config.output_dir, "best_model_by_metric"))
            unwrapped_model = accelerator.unwrap_model(ppo_trainer.model)
            unwrapped_model.llm.save_pretrained(os.path.join(config.output_dir, "best_model_by_metric/llm_lora_adapter"))
            if hasattr(unwrapped_model, 'clip') and hasattr(config, 'use_lora_clip') and config.use_lora_clip: 
                unwrapped_model.clip.save_pretrained(os.path.join(config.output_dir, "best_model_by_metric/clip_lora_adapter"))
            if hasattr(unwrapped_model, 'clip2llm_proj'):
                torch.save(unwrapped_model.clip2llm_proj.state_dict(), os.path.join(config.output_dir, "best_model_by_metric/clip2llm_proj.pt"))
        else:
            early_stop_counter += 1
            accelerator.print(f"[INFO] 早停计数器: {early_stop_counter}/{config.early_stopping_patience}")
            if early_stop_counter >= config.early_stopping_patience:
                accelerator.print(f"[INFO] 早停触发：连续{config.early_stopping_patience}轮未提升，终止训练")
                break

        ppo_trainer.save_pretrained(os.path.join(config.output_dir, "last_checkpoint"))
        with open(os.path.join(config.output_dir, "last_checkpoint", "epoch.txt"), "w") as f:
            f.write(str(epoch))
        unwrapped_model = accelerator.unwrap_model(ppo_trainer.model)
        unwrapped_model.llm.save_pretrained(os.path.join(config.output_dir, "last_checkpoint/llm_lora_adapter"))
        if hasattr(unwrapped_model, 'clip') and hasattr(config, 'use_lora_clip') and config.use_lora_clip:
            unwrapped_model.clip.save_pretrained(os.path.join(config.output_dir, "last_checkpoint/clip_lora_adapter"))
        if hasattr(unwrapped_model, 'clip2llm_proj'):
            torch.save(unwrapped_model.clip2llm_proj.state_dict(), os.path.join(config.output_dir, "last_checkpoint/clip2llm_proj.pt"))


    accelerator.end_training()